from src.algo.image_smilarity.model import simlarity_model as model
from src.utils import image_utils
from src.utils import matrix_utils
from .model_implements.mobilenet_v3 import ModelnetV3
from .model_implements.vit_base import VitBase
from .model_implements.bit import BigTransfer
from tqdm import tqdm


class Similarity:
    def __init__(self,model_name = "Big Transfer (BiT)",image_type="array"):
        '''

        :param model_name: 图片向量化模型 ["Mobilenet V3","Big Transfer (BiT)","Vision Transformer"]
        '''
        self.image_type = image_type
        if model_name == "Big Transfer (BiT)":
            model_cls = BigTransfer()
        elif model_name == "Mobilenet V3":
            model_cls = ModelnetV3()
        elif model_name == "Vision Transformer":
            model_cls =  VitBase()
            self.image_type = "pil"
        else:
            raise  ValueError("model_name输入有误")


        self.model = model.SimilarityModel(name=model_name, image_size=224, model_cls=model_cls)


    def get_image_feature(self,image_urls,batch_size=100):
        imgs = []
        features = []
        for index, url in tqdm(enumerate(image_urls)):
            if url == "":
                continue
            imgs.append(image_utils.load_image_url(url, required_size=(self.model.image_size, self.model.image_size),
                                                  image_type=self.image_type))
            if index !=0 and index % batch_size == 0:
                features.extend(self.model.model_cls.extract_feature(imgs))
                imgs = []
        if imgs:
            features.extend(self.model.model_cls.extract_feature(imgs))
        return features

    def check_similarity(self, img_urls, model):
        imgs = []
        for url in img_urls:
            if url == "": continue
            imgs.append(image_utils.load_image_url(url, required_size=(model.image_size, model.image_size), image_type=model.image_input_type))
        
        features = model.model_cls.extract_feature(imgs)
        results = []
        for i, v in enumerate(features):
            if i == 0: continue 
            dist = matrix_utils.cosine(features[0], v)
            print(f'{i} -- distance: {dist}')
            # results.append((imgs[i], f'similarity: {int(dist*100)}%'))
            original_img = image_utils.load_image_url(img_urls[i], required_size=None, image_type='pil')
            results.append((original_img, f'similarity: {int(dist*100)}%'))
        return results

    